{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from scipy import stats\n",
    "from xgrads import *\n",
    "import numpy as np  # 调用numpy\n",
    "import xarray as xr\n",
    "from math import *\n",
    "import matplotlib.pyplot as plt\n",
    "import cartopy.crs as ccrs\n",
    "from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def createmap():  ###############################################生成地图##########################################################\n",
    "    box = [-180, 180, -90, 90]  # 经度维度\n",
    "    scale = '110m'  # 地图分辨率\n",
    "    xstep = 20  # 下面标注经纬度的步长\n",
    "    ystep = 10\n",
    "    proj = ccrs.PlateCarree(central_longitude=180)  # 确定地图投影\n",
    "    fig = plt.figure(figsize=(9, 6))  # dpi=150)###生成底图\n",
    "    ax = fig.subplots(1, 1, subplot_kw={'projection': proj})  # 确定子图，与grads的类似\n",
    "    ##海岸线\n",
    "    ax.coastlines(scale)\n",
    "    # 标注坐标轴\n",
    "    ax.set_xticks(np.arange(box[0], box[1] + xstep, xstep), crs=ccrs.PlateCarree())\n",
    "    ax.set_yticks(np.arange(box[2], box[3] + ystep, ystep), crs=ccrs.PlateCarree())\n",
    "    # 经纬度格式，把0经度设置不加E和W\n",
    "    lon_formatter = LongitudeFormatter(zero_direction_label=False)\n",
    "    lat_formatter = LatitudeFormatter()\n",
    "    ax.xaxis.set_major_formatter(lon_formatter)\n",
    "    ax.yaxis.set_major_formatter(lat_formatter)\n",
    "    ############################################################################################################\n",
    "    return ax, fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "############NCEP_Z200_30y_Wt转换成nc文件############\n",
    "hgt = open_CtlDataset('D:\\\\grads\\\\TongJi\\\\NCEP_Z200_30y_Wt.ctl')\n",
    "hgt.attrs['pdef'] = 'None'\n",
    "hgt.to_netcdf('D:\\\\grads\\\\TongJi\\\\NCEP_Z200_30y_Wt.nc')\n",
    "hgt = xr.open_dataset('D:\\\\grads\\\\TongJi\\\\NCEP_Z200_30y_Wt.nc')\n",
    "############NCEP_TPSST_30y_Wt转换成nc文件###########\n",
    "air = open_CtlDataset('D:\\\\grads\\\\TongJi\\\\NCEP_TPSST_30y_Wt.ctl')\n",
    "air.attrs['pdef'] = 'None'\n",
    "air.to_netcdf('D:\\\\grads\\\\TongJi\\\\NCEP_TPSST_30y_Wt.nc')\n",
    "air = xr.open_dataset('D:\\\\grads\\\\TongJi\\\\NCEP_TPSST_30y_Wt.nc')\n",
    "###########NCEP_IOSST_30y_Wt转换成nc文件###########\n",
    "airindia = open_CtlDataset('D:\\\\grads\\\\TongJi\\\\NCEP_IOSST_30y_Wt.ctl')\n",
    "airindia.attrs['pdef'] = 'None'\n",
    "airindia.to_netcdf('D:\\\\grads\\\\TongJi\\\\NCEP_IOSST_30y_Wt.nc')\n",
    "airindia = xr.open_dataset('D:\\\\grads\\\\TongJi\\\\NCEP_IOSST_30y_Wt.nc')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# 热带太平洋海温"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "stTP = air['st'][:, :, :]\n",
    "stIO = airindia['st'][:, :, :]\n",
    "hgt = hgt['hgt'][:, :, :]\n",
    "Time = hgt['time'][:]\n",
    "lon = hgt['lon'][:]\n",
    "lat = hgt['lat'][:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "stTP = stTP.where(stTP < 99999, np.NAN)\n",
    "stIO = stIO.where(stIO < 99999, np.NAN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "TPshijian = stTP.mean(dim='lon', skipna=True).mean(dim='lat', skipna=True)\n",
    "IOshijian = stIO.mean(dim='lon', skipna=True).mean(dim='lat', skipna=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "avehgt = hgt.mean(dim='time')\n",
    "aveTP = TPshijian.mean()\n",
    "aveIO = IOshijian.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def yiyuanhuigui(temp, ave, aveh, h, t):\n",
    "    tempxy = np.zeros((len(lat), len(lon)))\n",
    "    tempx2 = 0\n",
    "    for t_i in range(len(t)):\n",
    "        tempxy = tempxy + temp[t_i] * h[t_i, :, :]\n",
    "        tempx2 = tempx2 + temp[t_i] * temp[t_i]\n",
    "    b = (tempxy - len(t) * ave * aveh) / (tempx2 - len(t) * pow(ave, 2))\n",
    "    return b\n",
    "\n",
    "\n",
    "def yiyuanb0(b, ave, aveh):\n",
    "    b0 = aveh - b * ave\n",
    "    return b0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "lon1, lat1 = np.meshgrid(lon, lat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "##热带太平洋\n",
    "b = yiyuanhuigui(TPshijian, aveTP, avehgt, hgt, Time)\n",
    "ax, fig = createmap()\n",
    "colorbar = ax.contourf(lon1, lat1, b, cmap='bwr', transform=ccrs.PlateCarree())\n",
    "plt.colorbar(colorbar, extendrect='True', pad=0.03, fraction=0.04, shrink=1)\n",
    "ax.set_title('Z200andTPSST')\n",
    "plt.tight_layout()\n",
    "plt.savefig(r'D:\\grads\\TongJi\\ex4\\Z200andTPSST.png')\n",
    "plt.close()\n",
    "b0 = yiyuanb0(b, aveTP, avehgt)\n",
    "ax, fig = createmap()\n",
    "colorbar = ax.contourf(lon1, lat1, b0, cmap='bwr', transform=ccrs.PlateCarree())\n",
    "plt.colorbar(colorbar, extendrect='True', pad=0.03, fraction=0.04, shrink=1)\n",
    "ax.set_title('Z200andTPSSTb0')\n",
    "plt.tight_layout()\n",
    "plt.savefig(r'D:\\grads\\TongJi\\ex4\\Z200andTPSSTb0.png')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "##热带印度洋\n",
    "b = yiyuanhuigui(IOshijian, aveIO, avehgt, hgt, Time)\n",
    "ax, fig = createmap()\n",
    "colorbar = ax.contourf(lon1, lat1, b, cmap='bwr', transform=ccrs.PlateCarree())\n",
    "plt.colorbar(colorbar, extendrect='True', pad=0.03, fraction=0.04, shrink=1)\n",
    "ax.set_title('Z200andIOSST')\n",
    "plt.tight_layout()\n",
    "plt.savefig(r'D:\\grads\\TongJi\\ex4\\Z200andIOSST.png')\n",
    "plt.close()\n",
    "b0 = yiyuanb0(b, aveIO, avehgt)\n",
    "ax, fig = createmap()\n",
    "colorbar = ax.contourf(lon1, lat1, b0, cmap='bwr', transform=ccrs.PlateCarree())\n",
    "plt.colorbar(colorbar, extendrect='True', pad=0.03, fraction=0.04, shrink=1)\n",
    "ax.set_title('Z200andIOSSTb0')\n",
    "plt.tight_layout()\n",
    "plt.savefig(r'D:\\grads\\TongJi\\ex4\\Z200andIOSSTb0.png')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "a1 = np.array(TPshijian).reshape((30, 1))\n",
    "a2 = np.array(IOshijian).reshape((30, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "##多元回归\n",
    "a = np.full((len(Time), 1), 1)\n",
    "X = np.hstack((a, a1, a2))\n",
    "X1 = np.linalg.inv((X.T).dot(X)).dot(X.T)\n",
    "b = np.zeros((3, len(lat), len(lon)))\n",
    "for i in range(len(lon)):\n",
    "    for j in range(len(lat)):\n",
    "        b[:, j, i] = X1.dot(hgt[:, j, i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "##b1热带太平洋\n",
    "ax, fig = createmap()\n",
    "colorbar = ax.contourf(lon1, lat1, b[1, :, :], cmap='bwr', transform=ccrs.PlateCarree())\n",
    "plt.colorbar(colorbar, extendrect='True', pad=0.03, fraction=0.04, shrink=1)\n",
    "ax.set_title('Z200andTPSSTb1')\n",
    "plt.tight_layout()\n",
    "plt.savefig(r'D:\\grads\\TongJi\\ex4\\Z200andTPSSTb1.png')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "##b2热带印度洋\n",
    "ax, fig = createmap()\n",
    "colorbar = ax.contourf(lon1, lat1, b[2, :, :], cmap='bwr', transform=ccrs.PlateCarree())\n",
    "plt.colorbar(colorbar, extendrect='True', pad=0.03, fraction=0.04, shrink=1)\n",
    "ax.set_title('Z200andIOSSTb2')\n",
    "plt.tight_layout()\n",
    "plt.savefig(r'D:\\grads\\TongJi\\ex4\\Z200andIOSSTb2.png')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "##b0\n",
    "ax, fig = createmap()\n",
    "colorbar = ax.contourf(lon1, lat1, b[0, :, :], cmap='bwr', transform=ccrs.PlateCarree())\n",
    "plt.colorbar(colorbar, extendrect='True', pad=0.03, fraction=0.04, shrink=1)\n",
    "ax.set_title('b0')\n",
    "plt.tight_layout()\n",
    "plt.savefig(r'D:\\grads\\TongJi\\ex4\\b0.png')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#偏相关.\n",
    "def xiang_guan_xi_shu(ave, aveh, shijian, shijian2, lenT):\n",
    "    if aveh.size == 1:\n",
    "        temp1 = 0\n",
    "        temp2 = 0\n",
    "        temp3 = 0\n",
    "        for t_i in range(len(lenT)):\n",
    "            temp1 = temp1 + (shijian[t_i] - ave) * (shijian2[t_i] - aveh)\n",
    "            temp2 = temp2 + pow(shijian[t_i] - ave, 2)\n",
    "            temp3 = temp3 + pow(shijian2[t_i] - aveh, 2)\n",
    "        temp1 = temp1 / len(lenT)\n",
    "        temp2 = sqrt(temp2 / len(lenT))\n",
    "        temp3 = sqrt(temp3 / len(lenT))\n",
    "    else:\n",
    "        temp1 = np.zeros((len(lat), len(lon)))\n",
    "        temp2 = 0\n",
    "        temp3 = np.zeros((len(lat), len(lon)))\n",
    "        for t_i in range(len(lenT)):\n",
    "            temp1 = temp1 + (shijian[t_i] - ave) * (shijian2[t_i, :, :] - aveh[:, :])\n",
    "            temp2 = temp2 + pow(shijian[t_i] - ave, 2)\n",
    "            temp3 = temp3 + (shijian2[t_i, :, :] - aveh[:, :]) * (shijian2[t_i, :, :] - aveh[:, :])\n",
    "        temp1 = temp1 / len(lenT)\n",
    "        temp2 = sqrt(temp2 / len(lenT))\n",
    "        temp3 = np.sqrt(temp3 / len(lenT))\n",
    "    R = temp1 / (temp2 * temp3)\n",
    "\n",
    "    return R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "r12 = xiang_guan_xi_shu(aveTP, aveIO, TPshijian, IOshijian, Time)\n",
    "r1y = xiang_guan_xi_shu(aveTP, avehgt, TPshijian, hgt, Time)\n",
    "r2y = xiang_guan_xi_shu(aveIO, avehgt, IOshijian, hgt, Time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ry12 = (r1y - r12 * r2y) / np.sqrt((1 - r2y * r2y) * (1 - r12 * r12))\n",
    "ry21 = (r2y - r12 * r1y) / np.sqrt((1 - r1y * r1y) * (1 - r12 * r12))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "##ry12\n",
    "ax, fig = createmap()\n",
    "colorbar = ax.contourf(lon1, lat1, ry12, cmap='bwr', transform=ccrs.PlateCarree())\n",
    "plt.colorbar(colorbar, extendrect='True', pad=0.03, fraction=0.04, shrink=1)\n",
    "ax.set_title('ry12')\n",
    "plt.tight_layout()\n",
    "plt.savefig(r'D:\\grads\\TongJi\\ex4\\Ry12.png')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#ry21\n",
    "ax, fig = createmap()\n",
    "colorbar = ax.contourf(lon1, lat1, ry21, cmap='bwr', transform=ccrs.PlateCarree())\n",
    "plt.colorbar(colorbar, extendrect='True', pad=0.03, fraction=0.04, shrink=1)\n",
    "ax.set_title('ry21')\n",
    "plt.tight_layout()\n",
    "plt.savefig(r'D:\\grads\\TongJi\\ex4\\Ry21.png')\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
